package mq

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"time"

	"github.com/imdario/mergo"
	"github.com/pkg/errors"
	"github.com/redis/go-redis/v9"
	"github.com/spf13/cast"
	"github.com/ssgreg/repeat"
	"go.uber.org/atomic"
)

var defaultConsumerOptions = ConsumerOptions{
	Group:             "group",
	ShardsCount:       10,                     // 分片队列的个数
	PrefetchCount:     100,                    // 每次从分片队列取消息的个数
	Block:             time.Second,            // Millisecond is minimal for Redis
	PendingBufferSize: 1000,                   // 本地消息缓存队列大小，需要 >= ShardsCount*PrefetchCount
	PipeBufferSize:    100,                    // 每个分片队列 ack pipeline 的个数
	PipePeriod:        100 * time.Millisecond, // ack分片队列等待的最长时间
	Latency:           10,                     // 消息延迟处理超过一定秒数
}

type ConsumerOptions struct {
	Topic             string
	Group             string
	Consumer          string
	ShardsCount       int8
	PrefetchCount     int64
	Block             time.Duration
	PendingBufferSize int64
	PipeBufferSize    int64
	PipePeriod        time.Duration
	Latency           int64
	ErrorNotifier     Notifier
}

type Consumer struct {
	*client
	cAck     chan Message
	cCons    chan Message
	needStop bool
	notifier Notifier
	opt      ConsumerOptions
	stopped  bool
	wgAck    *sync.WaitGroup
	wgCons   *sync.WaitGroup
	rate     atomic.Int32 // 消费速率
	latency  atomic.Int64 // 消费延迟
}

func NewConsumer(ctx context.Context, opt ConsumerOptions, rdsCli redis.UniversalClient) (*Consumer, error) {
	if err := mergo.Merge(&opt, defaultConsumerOptions); err != nil {
		return nil, err
	}

	cli, err := newClient(ctx, opt.Topic, opt.Group, opt.ShardsCount, rdsCli)
	if err != nil {
		return nil, err
	}

	cn := &Consumer{
		cAck:     make(chan Message, opt.PrefetchCount*opt.PrefetchCount*2),
		cCons:    make(chan Message, opt.PrefetchCount*opt.PrefetchCount),
		client:   cli,
		notifier: opt.ErrorNotifier,
		opt:      opt,
		wgAck:    &sync.WaitGroup{},
		wgCons:   &sync.WaitGroup{},
	}

	cn.wgAck.Add(1)

	go cn.ack(ctx)
	go cn.showState()

	return cn, nil
}

func (c *Consumer) showState() {
	for range time.NewTicker(time.Second).C {
		if c.notifier != nil {
			c.notifier.AmiState(&State{
				Rate:    c.rate.Load(),
				Latency: int32(c.latency.Load()),
			})
		}
		c.rate.Store(0)
	}
}

func (c *Consumer) Start(ctx context.Context) chan Message {
	for i := 0; i < int(c.opt.ShardsCount); i++ {
		c.wgCons.Add(1)

		go c.consume(ctx, i)
	}

	return c.cCons
}

func (c *Consumer) Stop() {
	c.needStop = true

	c.wgCons.Wait()
	close(c.cCons)
	c.stopped = true
}

func (c *Consumer) Close() {
	close(c.cAck)
	c.wgAck.Wait()
}

func (c *Consumer) consume(ctx context.Context, shard int) {
	group := makeGroupName(c.opt.Topic, c.opt.Group)
	topic := makeTopicName(c.opt.Topic, shard)

	lastID := "0-0"
	checkBacklog := true

	for {
		if c.needStop {
			break
		}

		var id string
		if checkBacklog {
			id = lastID
		} else {
			id = ">"
		}

		var res []redis.XStream

		err := repeat.Repeat(
			repeat.Fn(func() error {
				var err error

				res, err = c.rds.XReadGroup(ctx, &redis.XReadGroupArgs{
					Block:    c.opt.Block,
					Consumer: c.opt.Consumer,
					Count:    c.opt.PrefetchCount,
					Group:    group,
					Streams:  []string{topic, id},
					NoAck:    false,
				}).Result()

				if err != nil && !errors.Is(err, redis.Nil) {
					if c.notifier != nil {
						c.notifier.AmiError(err)
					}

					return repeat.HintTemporary(err)
				}

				return nil
			}),
			repeat.StopOnSuccess(),
			repeat.WithDelay(repeat.FullJitterBackoff(500*time.Millisecond).Set()),
		)

		if err != nil {
			if c.notifier != nil {
				c.notifier.AmiError(err)
			}

			continue
		}

		if checkBacklog && (len(res) == 0 || len(res[0].Messages) == 0) {
			checkBacklog = false
			continue
		}

		for _, s := range res {
			for idx, m := range s.Messages {
				lastID = m.ID

				msg := Message{
					Group:  group,
					ID:     m.ID,
					Stream: topic,
				}

				v, ok := m.Values[dataField]
				if !ok {
					if c.notifier != nil {
						c.notifier.AmiError(errors.Errorf(`Incorrect message format: no "data" field in message with id %v`, m.ID))
					}

					c.Ack(msg)
					continue
				}

				m, ok := v.(string)
				if !ok {
					c.Ack(msg)
					continue
				}

				msg.Body = &Package{}
				if err = json.Unmarshal([]byte(m), &msg.Body); err != nil {
					c.Ack(msg)
					continue
				}

				c.cCons <- msg
				if idx > 2 || c.notifier == nil {
					continue
				}

				// check message latency
				if seq := strings.Split(msg.ID, "-"); len(seq) == 2 {
					latency := time.Now().UnixMilli() - cast.ToInt64(seq[0])
					if latency < 0 {
						latency = 0
					}
					c.latency.Store(latency)
					if latency > c.opt.Latency*1000 {
						c.notifier.AmiError(errors.Errorf("msg too long to delay, mid: %v, lag: %vms", msg.ID, latency))
					}
				}
			}
		}
	}

	c.wgCons.Done()
}

func (c *Consumer) Ack(m Message) {
	c.cAck <- m
	c.rate.Inc()
}

func (c *Consumer) ack(ctx context.Context) {
	started := time.Now()
	tick := time.NewTicker(c.opt.PipePeriod)

	toAck := make(map[string][]Message)
	cnt := make(map[string]int)

	for {
		var (
			doStop bool
			stream string
		)

		select {
		case m, more := <-c.cAck:
			if !more {
				doStop = true
			} else {
				stream = m.Stream

				if toAck[stream] == nil {
					toAck[stream] = make([]Message, c.opt.PipeBufferSize)
					cnt[stream] = 0
				}

				toAck[stream][cnt[stream]] = m
				cnt[stream]++
			}
		case <-tick.C:
		}

		if doStop {
			c.sendAckAllStreams(ctx, toAck, cnt)
			break
		}

		if cnt[stream] >= int(c.opt.PipeBufferSize) {
			c.sendAckStreamWithLock(ctx, toAck[stream][0:cnt[stream]])
			cnt[stream] = 0
		} else if time.Since(started) >= c.opt.PipePeriod && len(c.cAck) == 0 {
			c.sendAckAllStreams(ctx, toAck, cnt)
			started = time.Now()
		}
	}

	c.wgAck.Done()
}

func (c *Consumer) sendAckAllStreams(ctx context.Context, toAck map[string][]Message, cnt map[string]int) {
	for stream := range toAck {
		c.sendAckStreamWithLock(ctx, toAck[stream][0:cnt[stream]])
		cnt[stream] = 0
	}
}

func (c *Consumer) sendAckStreamWithLock(ctx context.Context, lst []Message) {
	if len(lst) == 0 {
		return
	}

	ids := make([]string, len(lst))
	for i, m := range lst {
		ids[i] = m.ID
	}

	c.wgAck.Add(1)

	go func() {
		c.sendAckStream(ctx, lst[0].Stream, lst[0].Group, ids)
		c.wgAck.Done()
	}()
}

func (c *Consumer) sendAckStream(ctx context.Context, topic string, group string, ids []string) {
	err := repeat.Repeat(
		repeat.Fn(func() error {
			pipe := c.rds.TxPipeline()
			pipe.XAck(ctx, topic, group, ids...)
			_, err := pipe.Exec(ctx)

			if err != nil {
				if c.notifier != nil {
					c.notifier.AmiError(err)
				}

				return repeat.HintTemporary(err)
			}

			return nil
		}),
		repeat.StopOnSuccess(),
		repeat.WithDelay(repeat.FullJitterBackoff(500*time.Millisecond).Set()),
	)

	if err != nil && c.notifier != nil {
		c.notifier.AmiError(err)
	}
}
